from sensor_msgs.msg import LaserScan
import rospy
import numpy as np
import os
import os.path as osp
from nav_msgs.msg import Odometry
from std_msgs.msg import String
import tf
from tf.transformations import quaternion_from_euler

FILE_PATH = "/data/OGM-datasets/OGM-Turtlebot2/train"
SPLIT = "train"

class Dataset(object):
    def __init__(self, file_path, split):
        self.scan_file_names = []
        self.vel_file_names = []
        self.pos_file_names = []
        self.scans = []
        self.vels = []
        self.pos = []
        # open train.txt or dev.txt:
        scan_file_path = osp.join(osp.abspath(file_path), 'scans', split + '.txt')
        pos_file_path = osp.join(osp.abspath(file_path), 'positions', split + '.txt')
        vel_file_path = osp.join(osp.abspath(file_path), 'velocities', split + '.txt')
        with open(scan_file_path, 'r') as f:
            for line in f.read().split("\n"):
                if '.npy' in line:
                    self.scan_file_names.append(osp.join(osp.abspath(file_path), 'scans', line))
                    self.scans.append(np.load(osp.join(osp.abspath(file_path), 'scans', line)))
        with open(pos_file_path, 'r') as f:
            for line in f.read().split("\n"):
                if '.npy' in line:
                    self.pos_file_names.append(osp.join(osp.abspath(file_path), 'positions', line))
                    self.pos.append(np.load(osp.join(osp.abspath(file_path), 'positions', line)))
        with open(vel_file_path, 'r') as f:
            for line in f.read().split("\n"):
                if '.npy' in line:
                    self.vel_file_names.append(osp.join(osp.abspath(file_path), 'velocities', line))
                    self.vels.append(np.load(osp.join(osp.abspath(file_path), 'velocities', line)))
        self.length = len(self.scan_file_names)
        print("dataset length: ", self.length)

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        assert idx <= self.length
        data = {}
        data["scan"] = self.scans[idx]
        data["vels"] = self.vels[idx]
        data["pos"] = self.pos[idx]
        return data


if __name__ == "__main__":
    scan_file_path = osp.join(osp.abspath("/data/OGM-datasets/OGM-Turtlebot2/train"), 'scans', "train" + '.txt')
    pos_file_path = osp.join(osp.abspath("/data/OGM-datasets/OGM-Turtlebot2/train"), 'positions', "train" + '.txt')
    vel_file_path = osp.join(osp.abspath("/data/OGM-datasets/OGM-Turtlebot2/train"), 'velocities', "train" + '.txt')

    scan_file_names = []
    pos_file_names = []
    vel_file_names = []

    with open(scan_file_path, 'r') as f:
        for line in f.read().split("\n"):
            if '.npy' in line:
                scan_file_names.append(osp.join(osp.abspath("/data/OGM-datasets/OGM-Turtlebot2/train"), 'scans', line))
    with open(pos_file_path, 'r') as f:
        for line in f.read().split("\n"):
            if '.npy' in line:
                pos_file_names.append(osp.join(osp.abspath("/data/OGM-datasets/OGM-Turtlebot2/train"), 'positions', line))
    with open(vel_file_path, 'r') as f:
        for line in f.read().split("\n"):
            if '.npy' in line:
                vel_file_names.append(osp.join(osp.abspath("/data/OGM-datasets/OGM-Turtlebot2/train"), 'velocities', line))

    seq = 0
    rospy.init_node("raw_sensor_viz", anonymous=False )
    rate = rospy.Rate(10)
    scan_publisher = rospy.Publisher("scan", LaserScan, queue_size=10)
    odom_publisher = rospy.Publisher("odom", Odometry, queue_size=10)
    while not rospy.is_shutdown():
        # Synthesize scan message
        scan = LaserScan()
        scan.header.stamp = rospy.Time.now()
        scan.header.frame_id = "odom"
        scan.header.seq = seq
        ranges = np.load(scan_file_names[seq])
        scan.range_max = 100
        scan.range_min = 0.1
        scan.angle_min = (-135 * np.pi/180)
        scan.angle_max = (135 * np.pi/180)
        scan.angle_increment = (270 * np.pi/180) / 1080
        scan.ranges = ranges
        scan_publisher.publish(scan)

        # Synthesize Odom message
        pos = np.load(pos_file_names[seq])

        odom = Odometry()
        odom.header.stamp = rospy.Time.now()
        odom.header.frame_id = "odom"
        scan.header.seq = seq
        odom.pose.pose.position.x = pos[0]
        odom.pose.pose.position.y = pos[1]
        odom.pose.pose.position.z = 0.
        quat = quaternion_from_euler(0, 0, pos.tolist()[-1])
        odom.pose.pose.orientation = quat.tolist()
        odom.pose.covariance = [1,0,0,0,0,0, 0,1,0,0,0,0, 0,0,1,0,0,0, 0,0,0,1,0,0, 0,0,0,0,1,0, 0,0,0,0,0,1]
        odom_publisher.publish(odom)

        # Synthesize tf data
        br = tf.TransformBroadcaster()
        br.sendTransform((pos[0], pos[1], 0.),
                         quaternion_from_euler(0, 0, pos.tolist()[-1]),
                         rospy.Time.now(),
                         "odom",
                         "map")


        seq += 1
        rate.sleep()


